#include "KSocket.h"
#include "KKey.h"
#include "KSockAddress.h"

KSocket::LPFN_CANCELIOEX KSocket::s_cancel = NULL;

KSocket::KSocket() :
    m_timeOut(0),
    m_key(NULL),
    m_accept(NULL),
    m_connect(NULL),
    m_disconnect(NULL),
    m_transmit(NULL),
    m_socket(INVALID_SOCKET)
{
}

KSocket::~KSocket()
{
    delete m_key;
    close();
}

bool KSocket::getLocalAddress(KIpAddress &addrLocal) const
{
    assert(isAlive());
    SOCKADDR_STORAGE name;
    int length = sizeof name;
    if (getsockname(m_socket, reinterpret_cast<SOCKADDR *>(&name), &length) != SOCKET_ERROR) {
        if (length >= sizeof(SOCKADDR_IN)) {
            addrLocal.reset(*reinterpret_cast<const SOCKADDR_IN *>(&name));
            return true;
        }
    }
    return false;
}

bool KSocket::getRemoteAddress(KIpAddress &addrRemote) const
{
    assert(isAlive());
    SOCKADDR_STORAGE name;
    int length = sizeof name;
    if (getpeername(m_socket, reinterpret_cast<SOCKADDR *>(&name), &length) != SOCKET_ERROR) {
        if (length >= sizeof(SOCKADDR_IN)) {
            addrRemote.reset(*reinterpret_cast<const SOCKADDR_IN *>(&name));
            return true;
        }
    }
    return false;
}

bool KSocket::updateAcceptContextWith(const KSocket &acceptor)
{
    assert(isAlive());
    SOCKET sock = acceptor.handle();
    return setsockopt(m_socket, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT,
        reinterpret_cast<char *>(&sock), sizeof sock) != SOCKET_ERROR;
}

bool KSocket::updateConnectContext()
{
    assert(isAlive());
    return setsockopt(m_socket, SOL_SOCKET, SO_UPDATE_CONNECT_CONTEXT, NULL, 0) != SOCKET_ERROR;
}

SOCKET KSocket::attach(SOCKET sock)
{
    SOCKET s = m_socket;
    m_socket = sock;
    return s;
}

bool KSocket::bind(const KIAddress &addrLocal)
{
    assert(isAlive() && !addrLocal.isNull());
    return ::bind(m_socket, addrLocal.name(), addrLocal.length()) != SOCKET_ERROR;
}

void KSocket::close()
{
    if (isAlive()) {
        if (closesocket(m_socket) == SOCKET_ERROR) {
            KDebugLib("closesocket() failed: %d", WSAGetLastError());
        }
        m_socket = INVALID_SOCKET;
    }
}

bool KSocket::open(int sendSize)
{
    bool opened = ((m_socket = WSASocket(AF_INET, sockType(), protocol(), NULL, 0, WSA_FLAG_OVERLAPPED)) != INVALID_SOCKET);
    if (opened && sendSize >= 0) {
        return setSendBufSize(sendSize);
    }
    return opened;
}

bool KSocket::shutdown(int how)
{
    return isAlive() ? !::shutdown(m_socket, how) : true;
}

bool KSocket::setBroadcast(bool enabled)
{
    assert(isAlive());
    DWORD val = enabled ? TRUE : FALSE;
    return setsockopt(m_socket, SOL_SOCKET, SO_BROADCAST, reinterpret_cast<char *>(&val), sizeof val) != SOCKET_ERROR;
}

int KSocket::connectTime() const
{
    assert(isAlive());
    int val = -1, len = sizeof val;
    getsockopt(m_socket, SOL_SOCKET, SO_CONNECT_TIME, reinterpret_cast<char *>(&val), &len);
    return val;
}

bool KSocket::setKeepAlive(bool enabled, ULONG time, ULONG interval)
{
    assert(isAlive());
    DWORD val = enabled ? TRUE : FALSE;
    if (setsockopt(m_socket, SOL_SOCKET, SO_KEEPALIVE, reinterpret_cast<char *>(&val), sizeof val) == SOCKET_ERROR) {
        return false;
    } else if (enabled) {
        tcp_keepalive inBuf = {0}, outBuf = {0};
        inBuf.onoff             = val;
        inBuf.keepalivetime     = time;
        inBuf.keepaliveinterval = interval;
        if (WSAIoctl(m_socket, SIO_KEEPALIVE_VALS, &inBuf, sizeof inBuf,
            &outBuf, sizeof outBuf, &val, NULL, NULL) == SOCKET_ERROR) {
            return false;
        }
    }
    return true;
}

bool KSocket::noDelay() const
{
    assert(isAlive());
    DWORD val = FALSE;
    int len = sizeof val;
    getsockopt(m_socket, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&val), &len);
    return !!val;
}

bool KSocket::setNoDelay(bool enabled)
{
    if (isAlive()) {
        DWORD val = enabled ? TRUE : FALSE;
        return setsockopt(m_socket, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<char *>(&val), sizeof val) != SOCKET_ERROR;
    } else {
        return false;
    }
}

bool KSocket::setNonblock(bool enabled)
{
    assert(isAlive());
    u_long mode = enabled;
    return ioctlsocket(m_socket, FIONBIO, &mode) != SOCKET_ERROR;
}

bool KSocket::setReuseAddr(bool enabled)
{
    assert(isAlive());
    DWORD val = enabled ? TRUE : FALSE;
    return setsockopt(m_socket, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<char *>(&val), sizeof val) != SOCKET_ERROR;
}

int KSocket::recvBufSize() const
{
    assert(isAlive());
    int val = -1, len = sizeof val;
    getsockopt(m_socket, SOL_SOCKET, SO_RCVBUF, reinterpret_cast<char *>(&val), &len);
    return val;
}

bool KSocket::setRecvBufSize(DWORD size)
{
    assert(isAlive());
    return setsockopt(m_socket, SOL_SOCKET, SO_RCVBUF, reinterpret_cast<char *>(&size), sizeof size) != SOCKET_ERROR;
}

int KSocket::sendBufSize() const
{
    assert(isAlive());
    int val = -1, len = sizeof val;
    getsockopt(m_socket, SOL_SOCKET, SO_SNDBUF, reinterpret_cast<char *>(&val), &len);
    return val;
}

bool KSocket::setSendBufSize(DWORD size)
{
    assert(isAlive());
    return setsockopt(m_socket, SOL_SOCKET, SO_SNDBUF, reinterpret_cast<char *>(&size), sizeof size) != SOCKET_ERROR;
}

bool KSocket::initCancel()
{
    if (KSocket::s_cancel) {
        return true;
    }
    OSVERSIONINFO osvi = {sizeof(OSVERSIONINFO)};
    // Windows Vista or higher version
    if (GetVersionEx(&osvi) && osvi.dwMajorVersion >= 6) {
        HMODULE hKernel32 = GetModuleHandle(TEXT("kernel32"));
        KSocket::s_cancel = reinterpret_cast<KSocket::LPFN_CANCELIOEX>
            (GetProcAddress(hKernel32, "CancelIoEx"));
        return KSocket::s_cancel != NULL;
    }
    return false;
}

bool KSocket::isTcpSocket(const KSocket &sock)
{
    return sock.protocol() == IPPROTO_TCP && sock.sockType() == SOCK_STREAM;
}

bool KSocket::isUdpSocket(const KSocket &sock)
{
    return sock.protocol() == IPPROTO_UDP && sock.sockType() == SOCK_DGRAM;
}

bool KSocket::isUnspecific(const KSocket &sock)
{
    return !sock.protocol() && !sock.sockType();
}

bool KSocket::checkSetLinger(bool onoff, u_short sec, bool check)
{
    if (!check) {
        assert(isAlive());
    } else if (!isAlive()) {
        return true;
    }
    LINGER lin;
    lin.l_onoff = onoff;
    lin.l_linger = sec;
    return setsockopt(m_socket, SOL_SOCKET, SO_LINGER, reinterpret_cast<char *>(&lin), sizeof lin) != SOCKET_ERROR;
}

bool KSocket::initAcceptEx()
{
    if (!m_accept) {
        DWORD   bytes;
        GUID    guidAcceptEx = WSAID_ACCEPTEX;
        if (WSAIoctl(m_socket, SIO_GET_EXTENSION_FUNCTION_POINTER,
            &guidAcceptEx, sizeof guidAcceptEx,
            &m_accept, sizeof m_accept,
            &bytes, NULL, NULL
        ) == SOCKET_ERROR) {
            m_accept = NULL;
            return false;
        }
    }
    return true;
}

bool KSocket::initConnectEx()
{
    if (!m_connect) {
        DWORD   bytes;
        GUID    guidConnectEx = WSAID_CONNECTEX;
        if (WSAIoctl(m_socket, SIO_GET_EXTENSION_FUNCTION_POINTER,
            &guidConnectEx, sizeof guidConnectEx,
            &m_connect, sizeof m_connect,
            &bytes, NULL, NULL
        ) == SOCKET_ERROR) {
            m_connect = NULL;
            return false;
        }
    }
    return true;
}

bool KSocket::initDisconnectEx()
{
    if (!m_disconnect) {
        DWORD   bytes;
        GUID    guidDisconnectEx = WSAID_DISCONNECTEX;
        if (WSAIoctl(m_socket, SIO_GET_EXTENSION_FUNCTION_POINTER,
            &guidDisconnectEx, sizeof guidDisconnectEx,
            &m_disconnect, sizeof m_disconnect,
            &bytes, NULL, NULL
        ) == SOCKET_ERROR) {
            m_disconnect = NULL;
            return false;
        }
    }
    return true;
}

bool KSocket::initTransmitFile()
{
    if (!m_transmit) {
        DWORD   bytes;
        GUID    guidTransmitFile = WSAID_TRANSMITFILE;
        if (WSAIoctl(m_socket, SIO_GET_EXTENSION_FUNCTION_POINTER,
            &guidTransmitFile, sizeof guidTransmitFile,
            &m_transmit, sizeof m_transmit,
            &bytes, NULL, NULL
        ) == SOCKET_ERROR) {
            m_transmit = NULL;
            return false;
        }
    }
    return true;
}

BOOL KSocket::asyncAccept(KSocket *client, PVOID buf, LPOVERLAPPED request)
{
    DWORD bytes;
    return m_accept(m_socket, client->handle(), buf, 0,
        sizeof(SOCKADDR_STORAGE) + 16, sizeof(SOCKADDR_STORAGE) + 16, &bytes, request);
}

bool KSocket::asyncCancel(LPOVERLAPPED request)
{
    if (KSocket::s_cancel) {
        if (!KSocket::s_cancel(reinterpret_cast<HANDLE>(m_socket), request)) {
            KDebugLib("CancelIoEx() failed: %d", GetLastError());
        }
        return true;
    }
    return false;
}

BOOL KSocket::asyncConnect(const KIAddress &addr, LPOVERLAPPED request)
{
    assert(!addr.isNull());
    return m_connect(m_socket, addr.name(), addr.length(), NULL, 0, NULL, request);
}

BOOL KSocket::asyncDisconnect(LPOVERLAPPED request)
{
    return m_disconnect(m_socket, request, 0, 0);
}

bool KSocket::asyncRecv(LPWSABUF pBuf, LPOVERLAPPED request)
{
    DWORD flags = 0;
    return WSARecv(m_socket, pBuf, 1, NULL, &flags, request, NULL) != SOCKET_ERROR;
}

bool KSocket::asyncSend(LPWSABUF pBuf, LPOVERLAPPED request)
{
    return WSASend(m_socket, pBuf, 1, NULL, 0, request, NULL) != SOCKET_ERROR;
}

BOOL KSocket::asyncTransmit(HANDLE hFile, DWORD total, DWORD per, LPTRANSMIT_FILE_BUFFERS pBuf, LPOVERLAPPED request)
{
    return m_transmit(m_socket, hFile, total, per, request, pBuf, TF_USE_KERNEL_APC);
}